import  os
import numpy as np
import struct
import matplotlib.pyplot as plt
import  os
def showdigit(d,r=28,c=28):
    assert type(d) == type(np.array([],dtype=np.int8))
    d = d.reshape(r,c)
    plt.imshow(d,'gray')
    plt.show()
    plt.close()

class DIGITS:
    trainset = []
    testset = []
    trainlabel = []
    testlabel = []
    trainfile = 'res/train-images.idx3-ubyte'
    trainlabelfile = 'res/train-labels.idx1-ubyte'
    testlabelfile = 'res/t10k-labels.idx1-ubyte'
    testfile = 'res/t10k-images.idx3-ubyte'
    def getImgsFromFile(self,f,s=True)->'np.parray':
        dataset = []
        f = open(f,'rb')
        buf = f.read(16)
        magic,count,rows,cols = struct.unpack('>IIII',buf)
        pixels = rows*cols
        fmt = ">"+pixels.__str__()+'B'

        for i in range(count):
            buf = f.read(pixels)
            digit = struct.unpack(fmt,buf)
            digit = np.array(digit)
            dataset.append(digit)
            if s is True:
                showdigit(digit)
        f.close()
        return dataset

    def getImgsLabelFromFile(self,f,s=True)->'np.parray':
        dataset = []
        f = open(f,'rb')
        buf = f.read(8)
        magic,count = struct.unpack('>II',buf)
        fmt = '<B'
        for i in range(count):
            buf = f.read(1)
            (digit,) = struct.unpack(fmt,buf)

            dataset.append(digit)
            if s is True:
                showdigit(digit)
        f.close()
        return dataset

    def getImgsLabelSet(self,test=False):
        if test is False:
            p = 'res/data/trainimglabelset.npy'
            file = self.trainlabelfile
        else:
            p='res/data/testimglabelset.npy'
            file = self.testlabelfile
        if os.path.exists(p):
            return  np.load(p)
        else:
            self.trainlabel = self.getImgsLabelFromFile(file,s=False)
            self.trainlabel = np.array(self.trainlabel,dtype=np.uint8)
            if not os.path.exists('res/data'):
                os.mkdir('res/data')
            np.save(p,self.trainlabel)
            return self.trainlabel

    def getImgsSet(self,test=False):
        if test is False:
            p = 'res/data/trainimgset.npy'
            file = self.trainfile
        else:
            p = 'res/data/testimgset.npy'
            file = self.testfile
        if os.path.exists(p):
            return  np.load(p)
        else:
            imgset = self.getImgsFromFile(file,s=False)
            imgset = np.array(imgset)
            imgset = imgset>50
            imgset = np.array(imgset,np.bool)
            if not os.path.exists('res/data'):
                os.mkdir('res/data')
            np.save(p,imgset)
            return imgset

    def getSimpleData(self):
        trp = 'res/data/simpletraindata.npy'
        tep = 'res/data/simpletestdataTrain.npy'
        if os.path.exists(trp):
            trdata = np.load(trp)
        else :
            train = self.getImgsSet()
            datat = []
            for i in train:
                i = np.array(i, dtype=np.uint16)
                i = i.reshape(28, 28)
                i = sum(i)
                datat.append(i)
            datatt = np.array(datat)
            np.save(trp, datat)
            trdata = datat
        if os.path.exists(tep):
            tedata = np.load(tep)
        else:
            test = self.getImgsSet(True)
            datatt = []
            test = self.getImgsSet(True)
            for i in test:
                i = np.array(i,dtype=np.uint16)
                i = i.reshape(28,28)
                i = sum(i)
                datatt.append(i)
            datatt = np.array(datatt)
            np.save(tep,datatt)
            tedata = datatt
        return trdata,tedata

